# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
# ~~~~~~~ 利用DCT隐写技术解码图像中的隐藏信息~~~~~~~ #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

from skimage.io import imread, imsave
from skimage.color import rgb2ycbcr, ycbcr2rgb
# from skimage import img_as_float, img_as_ubyte
import numpy as np
import sys

from DataHidingDemo.dct_processing import dct_process_channel as dct
from DataHidingDemo.dct_processing import idct_process_channel as idct
from DataHidingDemo.extra import file_to_bin, bin_to_file

from PIL import PngImagePlugin

PngImagePlugin.MAX_TEXT_CHUNK = 500 * (1024 ** 2)

_quality = 70


# 输入：stego RGB图像
# 输出：保存隐藏文件
def dct_reveal(img, out_name='output_file.txt', N=8,coefficient=8):
    a, b = img.shape[:2]
    c = 3 if len(img.shape) == 3 else 1

    if c == 3:
        bitstream = []

        # 转换YCbCr颜色空间
        newimg = rgb2ycbcr(img)

        # 计算每个信道上的DCT系数
        Y = dct(newimg[..., 0], _quality,coefficient)
        Cb = dct(newimg[..., 1], _quality,coefficient)
        Cr = dct(newimg[..., 2], _quality,coefficient)
        newimg = np.dstack((Y, Cb, Cr))

        # 从每个NxN块中的DCT系数中检索LSB
        for i in range(0, a, 8):
            for j in range(0, b, 8):
                for k in range(c):
                    bitstream.append(str(int(abs(newimg[i + 1, j + 1, k])) % 2))

        # 将位流转换为文件
        return bin_to_file(bitstream, out_name)


def main():
    print("Enter the source path of the image to decode:")
    file = input()
    image = imread(file)
    dct_reveal(image)


if __name__ == '__main__':
    main()
